import datetime
from typing import List, Dict, Tuple, Optional, Union, Generator

from common.geometry.geom import Geom
from catalog.models.base_model import init_db, BaseItem
from catalog.models.base_stac_item import BaseSTACItem

from catalog.models.known_models import KNOWN_MODELS
from catalog.raster.raster_remote_service import SUPPORTED_REMOTE_COLLECTIONS

ALL_KNOWN_MODELS = list(KNOWN_MODELS.keys())
ALL_KNOWN_MODELS.extend(SUPPORTED_REMOTE_COLLECTIONS.keys())


def get_one_item(
    collection_id: str,
    uid: str,
    valid_until: Optional[datetime.datetime] = None,
    create_db_conn: bool = False,
) -> Optional[BaseItem]:
    """Get an item from a collection by the item's uid.

    Parameters
    ----------
    collection_id : str
        Data collection id for the data of interest, e.g "COPERNICUS/SENTINEL2/L2A"
    uid : str
        The unique id for the item of interest.
    valid_until : Optional[datetime.datetime], optional
        Retrieve item if it is valid at the given datetime,
        by default None to retrieve item that is valid at the time of calling this method.
    create_db_conn : bool, optional
        Whether to create a new database connection in query, by default True to create a new one.
        Notice it is an internal argument reserved for internal developers to re-use existing database connection.

    Returns
    -------
    Optional[BaseItem]
        The found item (could be a BaseSTACItem or a BaseVectorItem), or None if cannot be found.
    """
    model = KNOWN_MODELS[collection_id]
    if create_db_conn:
        model._meta.database = init_db()
        model._meta.database.connect()
    return model.get_one_item(uid, valid_until=valid_until)


def get_many_items(
    collection_id: str,
    uids: List[str],
    valid_until: Optional[datetime.datetime] = None,
    create_db_conn: bool = False,
) -> Generator[Optional[BaseItem], None, None]:
    """Get many items from given list of uids.

    Parameters
    ----------
    uids : List[str]
        A list of unique ids for the items to be retrieved.
    valid_until : Optional[datetime.datetime], optional
        Retrieve items that are valid at the given datetime,
        by default None to only get valid items at the time of calling this method.
    create_db_conn : bool, optional
        Whether to create a new database connection in query, by default True to create a new one.
        Notice it is an internal argument reserved for internal developers to re-use existing database connection.

    Yields
    ------
    Generator[Optional[BaseItem], None, None]
        A generator of retrieved items.
    """
    model = KNOWN_MODELS[collection_id]
    if create_db_conn:
        model._meta.database = init_db()
        model._meta.database.connect()
    return model.get_many_items(uids, valid_until=valid_until)


def query_many_items(
    collection_id: str,
    spatial: Optional[Geom] = None,
    temporal: Optional[
        Tuple[Optional[datetime.datetime], Optional[datetime.datetime]]
    ] = None,
    properties: Optional[Dict] = None,
    create_db_conn: bool = True,
) -> List[BaseSTACItem]:
    """
    Parameters
    ----------
    collection_id : str
        Data collection id for the data of interest, e.g "COPERNICUS/SENTINEL2/L2A"
    spatial : Optional[Geom], optional
        Spatial query as a `Geom` object that intersects with items of interest, by default None not to do spatial query
    temporal : Optional[Tuple[Optional[datetime.datetime], Optional[datetime.datetime]] ], optional
        Temporal query, as a tuple of start_datetime and end_datetime, by default None not to do temporal query
    properties : Optional[Dict], optional
        Property queries as dictionary, see each collection's model class for more details,
        by default None not to perform any property based query
    create_db_conn : bool, optional
        Whether to create a new database connection in query, by default True to create a new one.
        Notice it is an internal argument reserved for internal developers to re-use existing database connection.

    Returns
    -------
    List[BaseSTACItem]
        A list of `BaseSTACItem` objects for the metadata retrieved.
    """
    model = KNOWN_MODELS[collection_id]
    if create_db_conn:
        model._meta.database = init_db()
        model._meta.database.connect()
    if properties is None:
        items = model.query_many_items(spatial=spatial, temporal=temporal)
    else:
        items = model.query_many_items(spatial=spatial, temporal=temporal, **properties)
    items = [item for item in items if item is not None]
    return items


def list_collections(filters: Optional[Union[List[str], str]] = None) -> List[str]:
    """List collection ids matching a specified list of conditions

    NOTE
    ----
    Currently a super simple one, with `filters` applied in an `or` manner.
    We could upgrade in the future when we actually have a lot of collections.

    Examples
    --------
    >> import catalog
    >> catalog.meta.list_collections()  # -> list all collection ids
    >> # -> list all collection ids from NASA or COPERNICUS
    >> catalog.meta.list_collections(filters=["NASA", "COPERNICUS"])
    """
    if filters is None:
        return ALL_KNOWN_MODELS

    res = set()
    if isinstance(filters, str):
        filters = [filters]
    for key in ALL_KNOWN_MODELS:
        for condition in filters:
            if condition in key:
                res.add(key)
    return sorted(list(res))


def read_collection_info(
    collection_id: str, include_stats: bool = False, include_dbinfo: bool = False
) -> Dict:
    """Read basic information about a given collection

    Parameters
    ----------
    collection_id : str
        Data collection id for the data of interest, e.g "COPERNICUS/SENTINEL2/L2A".
    include_stats : bool
        Whether to include basic statistics for the collection, default is False not to include.
        If True, currently only include the total number of scenes in the result.
    include_dbinfo: bool
        Whether to include information about database table for the collection, default is False not to include.
        If True, currently only include the table name and field definitions in the result.

    Returns
    -------
    Dict
        A dictionary containing the collection's information.
        Following keys will be present: "collection_id", "default_resolutions", "bands" and "supported_vis".
        If `include_stats" is True, "total" will be present.
        If `include_dbinfo" is True, "table_name" and "fields" will be present.
    """
    if collection_id in KNOWN_MODELS:
        mdl = KNOWN_MODELS[collection_id]
        res = {"collection_id": mdl.collection_id}
        res["default_resolutions"] = mdl.default_resolutions
        res["bands"] = mdl.bands
        res["supported_vis"] = mdl.supported_vis

        if include_stats:
            res["total"] = mdl.count()

        if include_dbinfo:
            res["table_name"] = mdl._meta.table_name
            res["fields"] = mdl._meta.fields
    elif collection_id in SUPPORTED_REMOTE_COLLECTIONS:
        mdl = SUPPORTED_REMOTE_COLLECTIONS[collection_id]
        res = {"collection_id": collection_id}
        res["default_resolutions"] = mdl["default_resolutions"]
        res["bands"] = mdl["bands"]
        res["supported_vis"] = mdl["supported_vis"]
    return res
